package com.augrain.microframework.trace.context;

import com.augrain.microframework.core.spring.WebUtils;
import com.augrain.microframework.trace.common.TraceLabel;
import com.augrain.microframework.trace.common.spanid.SpanIdGenerator;
import com.augrain.microframework.trace.util.MDCUtils;
import com.augrain.microframework.trace.util.TraceUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.lang.Nullable;

/**
 * trace上下文
 *
 * @author biaoy
 * @since 2021/08/21
 */
@Slf4j
public class TraceContext {
    private static final ThreadLocal<TraceLabel> traceLabelTl = new InheritableThreadLocal<>();

    private TraceContext() {

    }

    public static void setTraceLabel(TraceLabel traceLabel) {
        traceLabelTl.set(traceLabel);
    }

    /**
     * 清空上下文
     */
    public static void clear() {
        traceLabelTl.remove();
        MDCUtils.clear();
    }

    /**
     * 快照, span_id 直接生成下一跳的span_id
     */
    @Nullable
    public static TraceContextSnapshot snapshot() {
        TraceLabel traceLabel = traceLabelTl.get();
        if (traceLabel != null) {
            TraceContextSnapshot snapshot = new TraceContextSnapshot();
            snapshot.setTraceId(traceLabel.getTraceId());
            snapshot.setSpanId(SpanIdGenerator.generateNextSpanId());
            snapshot.setSpanIdIndex(traceLabel.getSpanIdIndex());
            snapshot.setFromApp(traceLabel.getFromApp());
            snapshot.setFromIp(traceLabel.getFromIp());
            return snapshot;
        } else {
            return null;
        }
    }

    /**
     * 上下文继续
     */
    public static void continued(@Nullable TraceContextSnapshot snapshot) {
        TraceLabel traceLabel;
        if (snapshot != null) {
            traceLabel = TraceLabel.from(snapshot);
        } else {
            traceLabel = TraceUtils.init(WebUtils.getRequest());
        }
        setTraceLabel(traceLabel);

        // MDC状态恢复
        MDCUtils.continued(traceLabel);
    }

    public static String getTraceId() {
        TraceLabel traceLabel = getOrCreate();
        return traceLabel.getTraceId();
    }

    private static TraceLabel getOrCreate() {
        TraceLabel traceLabel = traceLabelTl.get();
        if (traceLabel == null) {
            traceLabel = TraceUtils.init(WebUtils.getRequest());
            setTraceLabel(traceLabel);
        }
        return traceLabel;
    }

    public static String getSpanId() {
        return traceLabelTl.get().getSpanId();
    }

    public static int getSpanIdIndex() {
        return traceLabelTl.get().getSpanIdIndex();
    }

    public static void updateSpanIdIndex(int spanIdIndex) {
        traceLabelTl.get().setSpanIdIndex(spanIdIndex);
    }
}
